{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import cv2\n",
    "import numpy as np\n",
    "import os\n",
    "import subprocess\n",
    "import glob\n",
    "from matplotlib import pyplot as plt\n",
    "\n",
    "INPUT_FOLDER = \"./input/\"\n",
    "OUTPUT_FOLDER = \"./output/\"\n",
    "TEST_IMAGE = \"test_image.jpg\"\n",
    "TEST_VIDEO = \"test_video.mov\"\n",
    "DEBUG = False\n",
    "\n",
    "def plot_image(image, title):\n",
    "    plt.imshow(image, cmap=plt.cm.gray)\n",
    "    plt.title(title)\n",
    "    plt.show()\n",
    "    \n",
    "def save_image(image, title):\n",
    "    cv2.imwrite(title, image)\n",
    "    \n",
    "def grayscale(image):\n",
    "    grayscale_image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)\n",
    "    if DEBUG:\n",
    "        plot_image(grayscale_image, \"grayscale\")\n",
    "    return grayscale_image\n",
    "\n",
    "def blur(image):\n",
    "    blur_image = cv2.GaussianBlur(image, (3,3),0)\n",
    "    if DEBUG:\n",
    "        plot_image(blur_image, \"blur\")\n",
    "    return blur_image\n",
    "        \n",
    "def canny(image):\n",
    "    canny_image = cv2.Canny(image, 100, 150)\n",
    "    if DEBUG:\n",
    "        plot_image(canny_image, \"canny\")\n",
    "    return canny_image\n",
    "        \n",
    "def roi(image):\n",
    "    bottom_padding = 100 # Front bumper compensation\n",
    "    height = image.shape[0]\n",
    "    width = image.shape[1]\n",
    "    # FYI, below values are highly dependant on the camera calibration i.e what part of the road is actually being captured\n",
    "    bottom_left = [0, height-bottom_padding]\n",
    "    bottom_right = [width, height-bottom_padding]\n",
    "    top_right = [width*1/3, height*1/3]\n",
    "    top_left = [width*2/3, height*1/3]\n",
    "    vertices = [np.array([bottom_left, bottom_right, top_left, top_right], dtype=np.int32)]\n",
    "    mask = np.zeros_like(image)  \n",
    "    cv2.fillPoly(mask, vertices, 255)\n",
    "    if DEBUG:\n",
    "        plot_image(mask, \"mask\")\n",
    "    masked_image = cv2.bitwise_and(image, mask)\n",
    "    if DEBUG:\n",
    "        plot_image(masked_image, \"roi\")\n",
    "    return masked_image\n",
    "    \n",
    "def averaged_lines(image, lines):\n",
    "    right_lines = []\n",
    "    left_lines = []\n",
    "    for x1,y1,x2,y2 in lines[:, 0]:\n",
    "        parameters = np.polyfit((x1, x2), (y1, y2), 1)\n",
    "        slope = parameters[0]\n",
    "        intercept = parameters[1]\n",
    "        if slope >= 0: \n",
    "            right_lines.append([slope, intercept])\n",
    "        else:\n",
    "            left_lines.append([slope, intercept])\n",
    "            \n",
    "    def merge_lines(image, lines):\n",
    "        if len(lines) > 0:\n",
    "            slope, intercept = np.average(lines, axis=0)\n",
    "            y1 = image.shape[0]\n",
    "            y2 = int(y1*(1/2))\n",
    "            x1 = int((y1 - intercept)/slope)\n",
    "            x2 = int((y2 - intercept)/slope)\n",
    "            return np.array([x1, y1, x2, y2])\n",
    "        \n",
    "    left = merge_lines(image, left_lines)\n",
    "    right = merge_lines(image, right_lines)\n",
    "    return left, right\n",
    "\n",
    "def hough_lines(image, rho, theta, threshold, min_line_len, max_line_gap):\n",
    "    lines_image = np.zeros((image.shape[0], image.shape[1], 3), dtype=np.uint8)\n",
    "    lines = cv2.HoughLinesP(image, rho, theta, threshold, np.array([]), minLineLength=min_line_len, maxLineGap=max_line_gap)\n",
    "    if lines is not None:\n",
    "        lines = averaged_lines(image, lines)\n",
    "        for line in lines:\n",
    "            if line is not None:\n",
    "                x1,y1,x2,y2 = line\n",
    "                cv2.line(lines_image, (x1, y1), (x2, y2), (0, 0, 255), 20)\n",
    "        if DEBUG:\n",
    "            plot_image(lines_image, \"lines\")\n",
    "    return lines_image\n",
    "\n",
    "def combine_images(image, initial_image, α=0.9, β=1.0, λ=0.0):\n",
    "    combined_image = cv2.addWeighted(initial_image, α, image, β, λ)\n",
    "    if DEBUG:\n",
    "        plot_image(combined_image, \"combined\")\n",
    "    return combined_image\n",
    "\n",
    "def find_street_lanes(image):\n",
    "    grayscale_image = grayscale(image)\n",
    "    blur_image = blur(grayscale_image)\n",
    "    canny_image = canny(blur_image)\n",
    "    roi_image = roi(canny_image)\n",
    "    hough_lines_image = hough_lines(roi_image, 0.9, np.pi/180, 100, 100, 50)\n",
    "    final_image = combine_images(hough_lines_image, image)\n",
    "    return final_image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Test Image\n",
    "test_image = cv2.imread(INPUT_FOLDER + TEST_IMAGE)\n",
    "street_lanes = find_street_lanes(test_image)\n",
    "save_image(street_lanes, OUTPUT_FOLDER + TEST_IMAGE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Test Video\n",
    "capture = cv2.VideoCapture(INPUT_FOLDER + TEST_VIDEO)\n",
    "while True:\n",
    "    _, frame = capture.read()\n",
    "    if frame is not None:\n",
    "        street_lanes = find_street_lanes(frame)\n",
    "        cv2.imshow(\"video\", street_lanes)\n",
    "        cv2.waitKey(20)\n",
    "    else:\n",
    "        break\n",
    "capture.release()\n",
    "cv2.destroyAllWindows()\n",
    "exit()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
